{
 "cells": [
  {
   "cell_type": "raw",
   "id": "df7d42b9-58a6-434c-a2d7-0b61142f6d3e",
   "metadata": {},
   "source": [
    "---\n",
    "sidebar_position: 4\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2195672-0cab-4967-ba8a-c6544635547d",
   "metadata": {},
   "source": [
    "# Handle Multiple Queries\n",
    "\n",
    "Sometimes, a query analysis technique may allow for multiple queries to be generated. In these cases, we need to remember to run all queries and then to combine the results. We will show a simple example (using mock data) of how to do that."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4079b57-4369-49c9-b2ad-c809b5408d7e",
   "metadata": {},
   "source": [
    "## Setup\n",
    "#### Install dependencies\n",
    "\n",
    "```{=mdx}\n",
    "import IntegrationInstallTooltip from \"@mdx_components/integration_install_tooltip.mdx\";\n",
    "import Npm2Yarn from \"@theme/Npm2Yarn\";\n",
    "\n",
    "<IntegrationInstallTooltip></IntegrationInstallTooltip>\n",
    "\n",
    "<Npm2Yarn>\n",
    "  @langchain/core @langchain/community @langchain/openai zod chromadb\n",
    "</Npm2Yarn>\n",
    "```\n",
    "\n",
    "#### Set environment variables\n",
    "\n",
    "We'll use OpenAI in this example:\n",
    "\n",
    "```\n",
    "OPENAI_API_KEY=your-api-key\n",
    "\n",
    "# Optional, use LangSmith for best-in-class observability\n",
    "LANGSMITH_API_KEY=your-api-key\n",
    "LANGCHAIN_TRACING_V2=true\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c20b48b8-16d7-4089-bc17-f2d240b3935a",
   "metadata": {},
   "source": [
    "### Create Index\n",
    "\n",
    "We will create a vectorstore over fake information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1f621694",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Module: null prototype] {\n",
       "  AdminClient: \u001b[36m[class AdminClient]\u001b[39m,\n",
       "  ChromaClient: \u001b[36m[class ChromaClient]\u001b[39m,\n",
       "  CloudClient: \u001b[36m[class CloudClient extends ChromaClient]\u001b[39m,\n",
       "  CohereEmbeddingFunction: \u001b[36m[class CohereEmbeddingFunction]\u001b[39m,\n",
       "  Collection: \u001b[36m[class Collection]\u001b[39m,\n",
       "  DefaultEmbeddingFunction: \u001b[36m[class _DefaultEmbeddingFunction]\u001b[39m,\n",
       "  GoogleGenerativeAiEmbeddingFunction: \u001b[36m[class _GoogleGenerativeAiEmbeddingFunction]\u001b[39m,\n",
       "  HuggingFaceEmbeddingServerFunction: \u001b[36m[class HuggingFaceEmbeddingServerFunction]\u001b[39m,\n",
       "  IncludeEnum: {\n",
       "    Documents: \u001b[32m\"documents\"\u001b[39m,\n",
       "    Embeddings: \u001b[32m\"embeddings\"\u001b[39m,\n",
       "    Metadatas: \u001b[32m\"metadatas\"\u001b[39m,\n",
       "    Distances: \u001b[32m\"distances\"\u001b[39m\n",
       "  },\n",
       "  JinaEmbeddingFunction: \u001b[36m[class JinaEmbeddingFunction]\u001b[39m,\n",
       "  OpenAIEmbeddingFunction: \u001b[36m[class _OpenAIEmbeddingFunction]\u001b[39m,\n",
       "  TransformersEmbeddingFunction: \u001b[36m[class _TransformersEmbeddingFunction]\u001b[39m\n",
       "}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import { Chroma } from \"@langchain/community/vectorstores/chroma\"\n",
    "import { OpenAIEmbeddings } from \"@langchain/openai\"\n",
    "import \"chromadb\";\n",
    "\n",
    "const texts = [\"Harrison worked at Kensho\", \"Ankush worked at Facebook\"]\n",
    "const embeddings = new OpenAIEmbeddings({ modelName: \"text-embedding-3-small\" })\n",
    "const vectorstore = await Chroma.fromTexts(\n",
    "    texts,\n",
    "    {},\n",
    "    embeddings,\n",
    "    {\n",
    "        collectionName: \"multi_query\"\n",
    "    }\n",
    ")\n",
    "const retriever = vectorstore.asRetriever(1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57396e23-c192-4d97-846b-5eacea4d6b8d",
   "metadata": {},
   "source": [
    "## Query analysis\n",
    "\n",
    "We will use function calling to structure the output. We will let it return multiple queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0b51dd76-820d-41a4-98c8-893f6fe0d1ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { z } from \"zod\";\n",
    "\n",
    "const searchSchema = z.object({\n",
    "    queries: z.array(z.string()).describe(\"Distinct queries to search for\")\n",
    "}).describe(\"Search over a database of job records.\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "783c03c3-8c72-4f88-9cf4-5829ce6745d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { ChatPromptTemplate } from \"@langchain/core/prompts\";\n",
    "import { RunnableSequence, RunnablePassthrough } from \"@langchain/core/runnables\";\n",
    "import { ChatOpenAI } from \"@langchain/openai\";\n",
    "\n",
    "const system = `You have the ability to issue search queries to get information to help answer user information.\n",
    "\n",
    "If you need to look up two distinct pieces of information, you are allowed to do that!`;\n",
    "\n",
    "const prompt = ChatPromptTemplate.fromMessages([\n",
    "    [\"system\", system],\n",
    "    [\"human\", \"{question}\"],\n",
    "])\n",
    "const llm = new ChatOpenAI({\n",
    "  modelName: \"gpt-3.5-turbo-0125\",\n",
    "  temperature: 0,\n",
    "});\n",
    "const llmWithTools = llm.withStructuredOutput(searchSchema, {\n",
    "  name: \"Search\"\n",
    "});\n",
    "const queryAnalyzer = RunnableSequence.from([\n",
    "  {\n",
    "      question: new RunnablePassthrough(),\n",
    "  },\n",
    "  prompt,\n",
    "  llmWithTools\n",
    "]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9564078",
   "metadata": {},
   "source": [
    "We can see that this allows for creating multiple queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc1d3863",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ queries: [ \u001b[32m\"Harrison\"\u001b[39m ] }"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await queryAnalyzer.invoke(\"where did Harrison Work\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af62af17-4f90-4dbd-a8b4-dfff51f1db95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ queries: [ \u001b[32m\"Harrison work\"\u001b[39m, \u001b[32m\"Ankush work\"\u001b[39m ] }"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await queryAnalyzer.invoke(\"where did Harrison and ankush Work\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7c65b2f-7881-45fc-a47b-a4eaaf48245f",
   "metadata": {},
   "source": [
    "## Retrieval with query analysis\n",
    "\n",
    "So how would we include this in a chain? One thing that will make this a lot easier is if we call our retriever asyncronously - this will let us loop over the queries and not get blocked on the response time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8dac7866",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { RunnableConfig, RunnableLambda } from \"@langchain/core/runnables\";\n",
    "\n",
    "const chain = async (question: string, config?: RunnableConfig) => {\n",
    "    const response = await queryAnalyzer.invoke(question, config);\n",
    "    const docs = [];\n",
    "    for (const query of response.queries) {\n",
    "        const newDocs = await retriever.invoke(query, config);\n",
    "        docs.push(...newDocs);\n",
    "    }\n",
    "    // You probably want to think about reranking or deduplicating documents here\n",
    "    // But that is a separate topic\n",
    "    return docs;\n",
    "}\n",
    "\n",
    "const customChain = new RunnableLambda({ func: chain });"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "232ad8a7-7990-4066-9228-d35a555f7293",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ Document { pageContent: \u001b[32m\"Harrison worked at Kensho\"\u001b[39m, metadata: {} } ]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await customChain.invoke(\"where did Harrison Work\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "28e14ba5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[\n",
       "  Document { pageContent: \u001b[32m\"Harrison worked at Kensho\"\u001b[39m, metadata: {} },\n",
       "  Document { pageContent: \u001b[32m\"Ankush worked at Facebook\"\u001b[39m, metadata: {} }\n",
       "]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await customChain.invoke(\"where did Harrison and ankush Work\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Deno",
   "language": "typescript",
   "name": "deno"
  },
  "language_info": {
   "file_extension": ".ts",
   "mimetype": "text/x.typescript",
   "name": "typescript",
   "nb_converter": "script",
   "pygments_lexer": "typescript",
   "version": "5.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
